#include "AccelerationStructure.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../RenderState.h"

AccelerationStructure::~AccelerationStructure()
{
	// vkDestroyAccelerationStructureKHR(GVulkanRHI->GetLogicalDevice(), mAccelerationStructure, nullptr);
	// vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), mAccelerationStructureBuffer, nullptr);
	// vkFreeMemory(GVulkanRHI->GetLogicalDevice(), mAccelerationStructureMemory, nullptr);

	GRenderState->DelayGCAccelerationStructure(mAccelerationStructure);
	GRenderState->DelayGCBuffer(mAccelerationStructureBuffer);
	GRenderState->DelayGCDeviceMemory(mAccelerationStructureMemory);
}

VkAccelerationStructureTypeKHR AccelerationStructure::GetType() const
{
	return mType;
}

VkBuildAccelerationStructureFlagsKHR AccelerationStructure::GetFlags() const
{
	return mFlags;
}

VkBuffer& AccelerationStructure::GetAccelerationStructureBuffer()
{
	return mAccelerationStructureBuffer;
}

VkDeviceMemory& AccelerationStructure::GetAccelerationStructureMemory()
{
	return mAccelerationStructureMemory;
}

VkAccelerationStructureKHR& AccelerationStructure::GetAccelerationStructure()
{
	return mAccelerationStructure;
}

void AccelerationStructure::CreateAccelerationStructureBuffer(VkDeviceSize size, VkBuffer& buffer, VkDeviceMemory& deviceMemory)
{
	VkBufferCreateInfo bufferCreateInfo{};
	bufferCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	bufferCreateInfo.size = size;
	bufferCreateInfo.usage = VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
	VK_CHECK_RESULT(vkCreateBuffer(GVulkanRHI->GetLogicalDevice(), &bufferCreateInfo, nullptr, &buffer));
	VkMemoryRequirements memoryRequirements{};
	vkGetBufferMemoryRequirements(GVulkanRHI->GetLogicalDevice(), buffer, &memoryRequirements);
	VkMemoryAllocateFlagsInfo memoryAllocateFlagsInfo{};
	memoryAllocateFlagsInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO;
	memoryAllocateFlagsInfo.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT_KHR;
	VkMemoryAllocateInfo memoryAllocateInfo{};
	memoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	memoryAllocateInfo.pNext = &memoryAllocateFlagsInfo;
	memoryAllocateInfo.allocationSize = memoryRequirements.size;
	memoryAllocateInfo.memoryTypeIndex = GVulkanRHI->FindMemoryType(memoryRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
	VK_CHECK_RESULT(vkAllocateMemory(GVulkanRHI->GetLogicalDevice(), &memoryAllocateInfo, nullptr, &deviceMemory));
	VK_CHECK_RESULT(vkBindBufferMemory(GVulkanRHI->GetLogicalDevice(), buffer, deviceMemory, 0));
}

void AccelerationStructure::CreateScratchBuffer(VkDeviceSize size, VkBuffer& buffer, VkDeviceMemory& deviceMemory, VkDeviceAddress& deviceAddress)
{
	VkBufferCreateInfo bufferCreateInfo{};
	bufferCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	bufferCreateInfo.size = size;
	bufferCreateInfo.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
	VK_CHECK_RESULT(vkCreateBuffer(GVulkanRHI->GetLogicalDevice(), &bufferCreateInfo, nullptr, &buffer));

	VkMemoryRequirements memoryRequirements{};
	vkGetBufferMemoryRequirements(GVulkanRHI->GetLogicalDevice(), buffer, &memoryRequirements);

	VkMemoryAllocateFlagsInfo memoryAllocateFlagsInfo{};
	memoryAllocateFlagsInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO;
	memoryAllocateFlagsInfo.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT_KHR;

	VkMemoryAllocateInfo memoryAllocateInfo{};
	memoryAllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	memoryAllocateInfo.pNext = &memoryAllocateFlagsInfo;
	memoryAllocateInfo.allocationSize = memoryRequirements.size;
	memoryAllocateInfo.memoryTypeIndex = GVulkanRHI->FindMemoryType(memoryRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
	VK_CHECK_RESULT(vkAllocateMemory(GVulkanRHI->GetLogicalDevice(), &memoryAllocateInfo, nullptr, &deviceMemory));
	VK_CHECK_RESULT(vkBindBufferMemory(GVulkanRHI->GetLogicalDevice(), buffer, deviceMemory, 0));

	VkBufferDeviceAddressInfoKHR bufferDeviceAddressInfo{};
	bufferDeviceAddressInfo.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
	bufferDeviceAddressInfo.buffer = buffer;
	deviceAddress = vkGetBufferDeviceAddressKHR(GVulkanRHI->GetLogicalDevice(), &bufferDeviceAddressInfo);
}

